# -*- coding: utf-8 -*-
# created on 2018/05/15

from mathsolver.functions.base import *
from mathsolver.functions.base.base import new_latex
from sympy.abc import n


class DengChaBiJiChuAssume(BaseFunction):
    """
    设等差数列的首项为a_(1),公差为a_(d).
    设等比数列的首项为a_(1),公比为a_(q).
    """
    def solver(self, *args):
        sls = args[0]
        for ShuLie in sls:
            a_name = ShuLie.aName
            tong_xiang = sympify(a_name)(n)
            if ShuLie.type == "DengBi":
                self.label.add("等比数列类")
                flag = []
                shou_xiang = ShuLie.aValue
                if shou_xiang is not None:
                    flag.append(1)
                else:
                    shou_xiang = sympify(a_name)(1)
                    if ShuLie.aRange.is_Union:
                        self.known['inequations'].append([shou_xiang, '!=', 0])
                    else:
                        left, right, lopen, ropen = ShuLie.aRange.args
                        if right.is_infinite and not left.is_infinite:
                            if lopen:
                                self.known['inequations'].append([shou_xiang, '>', left])
                            else:
                                self.known['inequations'].append([shou_xiang, '>=', left])
                    ShuLie.aValue = shou_xiang
                    flag.append(0)
                gongbi = ShuLie.qValue
                if gongbi is not None:
                    flag.append(1)
                else:
                    gongbi = sympify(a_name)(sympify("q"))
                    ShuLie.qValue = gongbi
                    flag.append(0)
                if isinstance(ShuLie.qRange, range):
                    pass
                elif ShuLie.qRange.is_Union:
                    ineqs = [[gongbi, '!=', 0]]
                    # self.known['inequations'].append([GongBi, '!=', 0])
                    for item in ShuLie.qRange.args:
                        left, right, lopen, ropen = item.args
                        if left.is_infinite and not right.is_infinite:
                            if ropen:
                                ineqs.append([gongbi, '<', right])
                            else:
                                ineqs.append([gongbi, '<=', right])
                        elif not left.is_infinite and not right.is_infinite:
                            if lopen and ropen:
                                ineqs.append([gongbi, '>', left])
                                ineqs.append([gongbi, '<', right])
                        elif not left.is_infinite and right.is_infinite:
                            if lopen:
                                ineqs.append([gongbi, '>', left])
                    if len(ineqs) > 2:
                        _ineqs = []
                        used_ineqs = []
                        for item in itertools.combinations(ineqs, 2):
                            if (item[0][1] == '<' and item[1][1] == '>') or (item[0][1] == '>' and item[1][1] == '<'):
                                if item[0][0] == item[1][0] and item[0][2] == item[1][2]:
                                    used_ineqs.append(item[0])
                                    used_ineqs.append(item[1])
                                    _ineqs.append([item[0][0], '!=', item[0][2]])
                        if len(_ineqs) > 0:
                            for _ineq in _ineqs:
                                self.known['inequations'].append(_ineq)
                            for tt in ineqs:
                                if tt not in self.known['inequations'] and tt not in used_ineqs:
                                    self.known['inequations'].append(tt)
                        else:
                            for ineq in ineqs:
                                self.known['inequations'].append(ineq)
                    else:
                        for ineq in ineqs:
                            self.known['inequations'].append(ineq)
                else:
                    left, right, lopen, ropen = ShuLie.qRange.args
                    if right.is_infinite and not left.is_infinite:
                        if lopen:
                            self.known['inequations'].append([gongbi, '>', left])
                        else:
                            self.known['inequations'].append([gongbi, '>=', left])
                if flag == [1, 0]:
                    self.steps.append(["", "设等比数列\\{%s\\}的公比为%s, 则" % (new_latex(tong_xiang), new_latex(gongbi))])
                elif flag == [0, 1]:
                    self.steps.append(["", "设等比数列\\{%s\\}的首项为%s, 则" % (new_latex(tong_xiang), new_latex(shou_xiang))])
                elif flag == [0, 0]:
                    self.steps.append(["", "设等比数列\\{%s\\}的首项为%s, 公比为%s, 则" % (new_latex(tong_xiang), new_latex(shou_xiang), new_latex(gongbi))])
                else:
                    self.steps.append(["", ""])
            else:
                self.label.add("等差数列类")
                flag = []
                shou_xiang = ShuLie.aValue
                if shou_xiang is not None:
                    flag.append(1)
                else:
                    shou_xiang = sympify(a_name)(1)
                    ShuLie.aValue = shou_xiang
                    flag.append(0)
                if ShuLie.aRange.is_Union:
                    left, right, lopen, ropen = ShuLie.aRange.args[0].args
                    _left, _right, _lopen, _ropen = ShuLie.aRange.args[1].args
                    if right == _left and ropen and _lopen:
                        self.known['inequations'].append([shou_xiang, '!=', right])
                else:
                    left, right, lopen, ropen = ShuLie.aRange.args
                    if right.is_infinite and not left.is_infinite:
                        if lopen:
                            self.known['inequations'].append([shou_xiang, '>', left])
                        else:
                            self.known['inequations'].append([shou_xiang, '>=', left])
                gongcha = ShuLie.dValue
                if gongcha is None:
                    gongcha = sympify(a_name)(sympify("d"))
                    flag.append(0)
                    ShuLie.dValue = gongcha
                else:
                    flag.append(1)
                if isinstance(ShuLie.dRange, range):
                    pass
                elif ShuLie.dRange.is_Union:
                    left, right, lopen, ropen = ShuLie.dRange.args[0].args
                    _left, _right, _lopen, _ropen = ShuLie.dRange.args[1].args
                    if right == _left and ropen and _lopen:
                        self.known['inequations'].append([gongcha, '!=', right])
                else:
                    left, right, lopen, ropen = ShuLie.dRange.args
                    if right.is_infinite and not left.is_infinite:
                        if lopen:
                            self.known['inequations'].append([gongcha, '>', left])
                        else:
                            self.known['inequations'].append([gongcha, '>=', left])
                if flag == [1, 0]:
                    self.steps.append(["", "设等差数列\\{%s\\}的公差为%s, 则" % (new_latex(tong_xiang), new_latex(gongcha))])
                elif flag == [0, 1]:
                    self.steps.append(["", "设等差数列\\{%s\\}的首项为%s, 则" % (new_latex(tong_xiang), new_latex(shou_xiang))])
                elif flag == [0, 0]:
                    self.steps.append(["", "设等差数列\\{%s\\}的首项为%s, 公差为%s, 则" % (new_latex(tong_xiang), new_latex(shou_xiang), new_latex(gongcha))])
                else:
                    self.steps.append(["", ""])
            self.output.append(ShuLie)
        return self


class DengChaBiEqNormalization(BaseFunction):
    """
    用基本量法化简等差,等比类方程/不等式
    """
    def solver(self, *args):
        eqs = args[0]
        if self.known['inequations']:
            for ineq in self.known['inequations']:
                left, op, right = ineq
                new_ineq = [sympify(left), op, sympify(right)]
                if new_ineq not in eqs:
                    eqs.append(new_ineq)
        if self.known['equations']:
            for eq in self.known['equations']:
                if eq not in eqs:
                    eqs.append(eq)
        conditon_eqs = []
        for eq in eqs:
            if len(eq) == 2:
                expr = sympify(eq[0]) - sympify(eq[1])
            else:
                expr = sympify(eq[0]) - sympify(eq[2])
            p1 = r'[a-zA-Z]_\(.*?\)'
            m1 = re.findall(p1, str(expr))
            m1 = set(m1)
            m1 = list(m1)
            if m1:
                for item1 in m1:
                    new_item1 = sympify(item1)
                    new_item1_name = new_item1.func
                    new_item1_index = new_item1.args[0]
                    if new_item1_index == sympify('d') or new_item1_index == sympify('q'):
                        continue
                    if new_item1_index.free_symbols:
                        conditon_eqs.append([new_item1_index, ">", S.Zero])
                    sl = self.search(str(new_item1_name))
                    if sl:
                        sl_type = sl.type
                        p2 = r'[A-Z]_\(.*?\)'
                        m2 = re.findall(p2, item1)
                        m2 = set(m2)
                        m2 = list(m2)
                        if m2:
                            if sl_type == "DengCha":
                                new_item1_value = sl.aValue * new_item1_index + new_item1_index * (
                                        new_item1_index - 1) * sl.dValue / 2
                                expr = expr.subs({new_item1: new_item1_value})
                            elif sl_type == "DengBi":
                                new_item1_value = sl.aValue * (1 - sl.qValue ** new_item1_index) / (
                                        1 - sl.qValue)
                                expr = expr.subs({new_item1: new_item1_value})
                        else:
                            if sl_type == "DengCha":
                                new_item1_value = sl.aValue + (new_item1_index - 1) * sl.dValue
                                expr = expr.subs({new_item1: new_item1_value})
                            elif sl_type == "DengBi":
                                if new_item1_index == 1:
                                    new_item1_value = sl.aValue
                                else:
                                    new_item1_value = sl.aValue * sl.qValue ** (new_item1_index - 1)
                                expr = expr.subs({new_item1: new_item1_value})
            if len(eq) == 2:
                if expr != 0:
                    conditon_eqs.append([expr, S.Zero])
            else:
                conditon_eqs.append([expr, eq[1], S.Zero])
        if len(conditon_eqs) == 1:
            if len(conditon_eqs[0]) == 2:
                self.steps.append(["", "%s" % BaseEq(conditon_eqs[0]).printing()])
                self.output.append(BaseEq(conditon_eqs[0]))
            else:
                self.steps.append(["", "%s" % BaseIneq(conditon_eqs[0]).printing()])
                self.output.append(BaseIneq(conditon_eqs[0]))
        else:
            self.steps.append(["", "%s" % BaseIneqs(conditon_eqs).printing()])
            self.output.append(BaseIneqs(conditon_eqs))
        self.known['Sequations'] = conditon_eqs
        self.label.add("基本量法")
        return self


class DengChaBiExprNormalization(BaseFunction):
    """
    用基本量法化简等差,等比类代数式
    """
    def solver(self, *args):
        expr = args[0].sympify()
        p1 = r'[a-zA-Z]_\(.*?\)'
        m1 = re.findall(p1, str(expr))
        m1 = set(m1)
        m1 = list(m1)
        if m1:
            for item1 in m1:
                new_item1 = sympify(item1)
                new_item1_name = new_item1.func
                new_item1_index = new_item1.args[0]
                if new_item1_index == sympify('d') or new_item1_index == sympify('q'):
                    continue
                sl = self.search(str(new_item1_name))
                sl_type = sl.type
                p2 = r'[A-Z]_\(.*?\)'
                m2 = re.findall(p2, item1)
                m2 = set(m2)
                m2 = list(m2)
                if m2:
                    if sl_type == "DengCha":
                        new_item1_value = sl.aValue * new_item1_index + new_item1_index * (new_item1_index - 1) * sl.dValue / 2
                        expr = expr.subs({new_item1: new_item1_value})
                        self.label.add("等差数列的前n项和")
                    elif sl_type == "DengBi":
                        new_item1_value = sl.aValue * (1 - sl.qValue ** new_item1_index) / (1 - sl.qValue)
                        expr = expr.subs({new_item1: new_item1_value})
                        self.label.add("等比数列的前n项和")
                else:
                    if sl_type == "DengCha":
                        new_item1_value = sl.aValue + (new_item1_index - 1) * sl.dValue
                        expr = expr.subs({new_item1: new_item1_value})
                        self.label.add("等差数列的通项公式")
                    elif sl_type == "DengBi":
                        new_item1_value = sl.aValue * sl.qValue ** (new_item1_index - 1)
                        expr = expr.subs({new_item1: new_item1_value})
                        self.label.add("等比数列的通项公式")
        self.steps.append(["", "∴%s = %s" % (new_latex(args[0].sympify()), new_latex(expr))])
        self.output.append(BasePoly(expr))
        return self


class ShuLieSnEqNormalization(BaseFunction):
    """
    将S_(?.n)化为a_(?.n)
    """
    def solver(self, *args):
        eq = args[0].sympify()
        a_name = args[1]
        if isinstance(args[0], BaseEq):
            expr1 = (eq[0] - eq[1]).expand().simplify()
        elif isinstance(args[0], BaseIneq):
            left, op, right = eq
            expr1 = (left - right).expand().simplify()
        else:
            raise Exception("to do")
        p1 = r'[A-Z]_\(.*?\)'
        m1 = re.findall(p1, str(expr1))
        m1 = set(m1)
        m1 = list(m1)
        conditon_eqs = []
        if m1:
            for item1 in m1:
                new_item1 = sympify(item1)
                new_item1_index = new_item1.args[0]
                new_item1_value = 0
                for i in range(1, new_item1_index + 1):
                    new_item1_value += a_name(i)
                expr1 = expr1.subs({new_item1: new_item1_value})
            conditon_eqs.extend([expr1, S.Zero])
        self.output.append(BaseEq(conditon_eqs))
        return self


class ShuLieSnExprNormalization(BaseFunction):
    """
    将S_(?.n)化为a_(?.n)
    """
    def solver(self, *args):
        expr = args[0].sympify()
        a_name = args[1]
        p = r'[A-Z]_\(.*?\)'
        m = re.findall(p, str(expr))
        m = set(m)
        m = list(m)
        if m:
            for item in m:
                new_item = sympify(item)
                new_item_index = new_item.args[0]
                new_item_value = 0
                for i in range(1, new_item_index + 1):
                    new_item_value += a_name(i)
                expr = expr.subs({new_item: new_item_value})
        self.output.append(BasePoly(expr))
        return self


class ShuLieAnNormalization(BaseFunction):
    """
    a_(?.n)化为S_(?.n)
    """
    def solver(self, *args):
        eq = args[0].sympify()
        s_name = args[1]
        if isinstance(args[0], BaseEq):
            expr1 = (eq[0] - eq[1]).expand().simplify()
        elif isinstance(args[0], BaseIneq):
            left, op, right = eq
            expr1 = (left - right).expand().simplify()
        else:
            raise Exception("to do")
        p1 = r'[a-z]_\(.*?\)'
        m1 = re.findall(p1, str(expr1))
        m1 = set(m1)
        m1 = list(m1)
        conditon_eqs = []
        if m1:
            for item1 in m1:
                new_item1 = sympify(item1)
                new_item1_index = new_item1.args[0]
                new_item1_value = s_name(new_item1_index) - s_name(new_item1_index - 1)
                expr1 = expr1.subs({new_item1: new_item1_value})
            conditon_eqs.extend([expr1, S.Zero])
        self.output.append(BaseEq(conditon_eqs))
        return self


class ShuLieAnExprNormalization(BaseFunction):
    """
    a_(?.n)化为S_(?.n)
    """
    def solver(self, *args):
        expr = args[0].sympify()
        s_name = args[1]
        p = r'[a-z]_\(.*?\)'
        m = re.findall(p, str(expr))
        m = set(m)
        m = list(m)
        if m:
            for item in m:
                new_item = sympify(item)
                new_item_index = new_item.args[0]
                if new_item_index == 1:
                    new_item_value = s_name(new_item_index)
                else:
                    new_item_value = s_name(new_item_index) - s_name(new_item_index - 1)
                expr = expr.subs({new_item: new_item_value})
        self.output.append(BaseValue(expr))
        return self
